import pyautogui
import tkinter as tk
from PIL import ImageGrab
import numpy as np
import math
from collections import deque


class ColorDetector:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("智能屏幕取色器-环状区域采样版")

        # 获取屏幕尺寸并计算中心点
        self.screen_width, self.screen_height = pyautogui.size()
        self.center_x = self.screen_width // 2
        self.center_y = self.screen_height // 2

        # 采样参数
        self.ring_radius = 8  # 采样环形区域半径(像素)
        self.sample_points = 8  # 采样点数
        self.dynamic_tolerance = True  # 启用动态容差

        # 历史颜色缓存（用于动态容差）
        self.color_history = deque(maxlen=10)

        # 快捷键绑定
        self.root.bind('<Control-s>', lambda event: self.set_target_color())
        self.root.bind('<Control-S>', lambda event: self.set_target_color())

        # 状态变量
        self.target_color = None
        self.tolerance = 10
        self.detecting = False

        # GUI布局
        self.setup_ui()

    def setup_ui(self):
        # 状态显示
        self.status_label = tk.Label(self.root, text="未选择目标颜色\n(使用 Ctrl+S 快速锁定)",
                                     font=('Arial', 12))
        self.status_label.pack(pady=10)

        # 当前颜色显示
        current_frame = tk.Frame(self.root)
        current_frame.pack(pady=5)
        tk.Label(current_frame, text="当前颜色:").pack(side=tk.LEFT)
        self.current_frame = tk.Frame(current_frame, width=50, height=30, relief="groove", bg='white')
        self.current_frame.pack(side=tk.LEFT, padx=5)

        # 目标颜色显示
        target_frame = tk.Frame(self.root)
        target_frame.pack(pady=5)
        tk.Label(target_frame, text="目标颜色:").pack(side=tk.LEFT)
        self.target_frame = tk.Frame(target_frame, width=50, height=30, relief="groove", bg='white')
        self.target_frame.pack(side=tk.LEFT, padx=5)

        # 控制面板
        control_frame = tk.Frame(self.root)
        control_frame.pack(pady=10)

        # 采样半径设置
        tk.Label(control_frame, text="采样半径:").grid(row=0, column=0, sticky='e')
        self.radius_entry = tk.Entry(control_frame, width=5)
        self.radius_entry.insert(0, str(self.ring_radius))
        self.radius_entry.grid(row=0, column=1, sticky='w')

        tk.Label(control_frame, text="采样点数:").grid(row=0, column=2, sticky='e')
        self.points_entry = tk.Entry(control_frame, width=5)
        self.points_entry.insert(0, str(self.sample_points))
        self.points_entry.grid(row=0, column=3, sticky='w')

        # 容差设置
        tk.Label(control_frame, text="容差:").grid(row=1, column=0, sticky='e')
        self.tolerance_entry = tk.Entry(control_frame, width=5)
        self.tolerance_entry.insert(0, str(self.tolerance))
        self.tolerance_entry.grid(row=1, column=1, sticky='w')

        # 动态容差选项
        self.dynamic_var = tk.BooleanVar(value=self.dynamic_tolerance)
        self.dynamic_check = tk.Checkbutton(control_frame, text="动态容差",
                                            variable=self.dynamic_var,
                                            command=self.toggle_dynamic_tolerance)
        self.dynamic_check.grid(row=1, column=2, columnspan=2)

        # 控制按钮
        button_frame = tk.Frame(control_frame)
        button_frame.grid(row=2, column=0, columnspan=4, pady=5)

        tk.Button(button_frame, text="锁定当前颜色 (Ctrl+S)",
                  command=self.set_target_color).pack(side=tk.LEFT, padx=5)
        tk.Button(button_frame, text="开始检测",
                  command=self.toggle_detection).pack(side=tk.LEFT, padx=5)
        tk.Button(button_frame, text="停止检测",
                  command=self.stop_detection).pack(side=tk.LEFT, padx=5)

        # 信息区域
        self.info_frame = tk.Frame(self.root)
        self.info_frame.pack(fill=tk.X, padx=10, pady=5)

        self.info_label = tk.Label(self.info_frame,
                                   text="等待启动... | Ctrl+S 锁定颜色 | Ctrl+C 退出",
                                   font=('Arial', 10))
        self.info_label.pack(side=tk.LEFT)

        self.dynamic_status = tk.Label(self.info_frame, text="动态容差: 未启用", fg="gray")
        self.dynamic_status.pack(side=tk.RIGHT)

    def toggle_dynamic_tolerance(self):
        self.dynamic_tolerance = self.dynamic_var.get()
        status = "启用" if self.dynamic_tolerance else "禁用"
        color = "green" if self.dynamic_tolerance else "gray"
        self.dynamic_status.config(text=f"动态容差: {status}", fg=color)

    def stop_detection(self):
        self.detecting = False
        self.status_label.config(text="检测已停止")
        self.info_label.config(text="检测已停止 | Ctrl+S 锁定颜色 | Ctrl+C 退出")
        self.root.config(bg="SystemButtonFace")

    def toggle_detection(self):
        """切换检测状态"""
        if not self.target_color:
            self.status_label.config(text="请先锁定目标颜色！")
            return

        self.detecting = not self.detecting
        if self.detecting:
            self.status_label.config(text="检测运行中")
            self.info_label.config(text="检测运行中 | Ctrl+C 退出")
            self.color_history.clear()  # 开始检测时清空历史记录
        else:
            self.stop_detection()

    def calculate_color_distance(self, color1, color2):
        """计算颜色差异（加权欧氏距离）"""
        # 使用加权距离（人眼对绿色更敏感）
        r_diff = color1[0] - color2[0]
        g_diff = color1[1] - color2[1]
        b_diff = color1[2] - color2[2]
        return math.sqrt(r_diff ** 2 * 0.3 + g_diff ** 2 * 0.59 + b_diff ** 2 * 0.11)

    def get_sample_points(self):
        """获取环形采样点坐标"""
        points = []
        self.ring_radius = int(self.radius_entry.get() or self.ring_radius)
        self.sample_points = int(self.points_entry.get() or self.sample_points)

        for i in range(self.sample_points):
            angle = 2 * math.pi * i / self.sample_points
            x = self.center_x + self.ring_radius * math.cos(angle)
            y = self.center_y + self.ring_radius * math.sin(angle)
            points.append((int(x), int(y)))
        return points

    def get_point_colors(self, points):
        """获取多个采样点的颜色"""
        colors = []
        try:
            # 只截取包含所有点的矩形区域
            min_x = min(p[0] for p in points)
            max_x = max(p[0] for p in points)
            min_y = min(p[1] for p in points)
            max_y = max(p[1] for p in points)

            # 截取包含所有采样点的区域
            bbox = (min_x, min_y, max_x + 1, max_y + 1)
            region = ImageGrab.grab(bbox=bbox)

            for point in points:
                # 转换为区域内的相对坐标
                rel_x = point[0] - min_x
                rel_y = point[1] - min_y

                # 确保坐标在图像范围内
                if 0 <= rel_x < region.width and 0 <= rel_y < region.height:
                    pixel = region.getpixel((rel_x, rel_y))
                    colors.append(pixel)
        except Exception as e:
            print(f"采样错误: {e}")

        return colors

    def calculate_average_color(self, points):
        """计算多个采样点的平均颜色"""
        colors = self.get_point_colors(points)
        if not colors:
            return (0, 0, 0)

        total_r, total_g, total_b = 0, 0, 0
        for color in colors:
            total_r += color[0]
            total_g += color[1]
            total_b += color[2]

        n = len(colors)
        return (
            int(total_r / n),
            int(total_g / n),
            int(total_b / n)
        )

    def update_tolerance(self):
        """动态调整容差阈值"""
        if not self.color_history or len(self.color_history) < 5:
            return self.tolerance

        # 计算最近几个样本的颜色差异平均值
        diffs = []
        for i in range(1, len(self.color_history)):
            diff = self.calculate_color_distance(self.color_history[i], self.color_history[i - 1])
            diffs.append(diff)

        if not diffs:
            return self.tolerance

        avg_diff = sum(diffs) / len(diffs)

        # 动态调整容差：当环境变化大时增大容差，环境稳定时减小容差
        dynamic_tolerance = int(self.tolerance * (1 + avg_diff / 30))
        return min(max(dynamic_tolerance, 5), 50)  # 保持在5-50之间

    def set_target_color(self, event=None):
        """设置目标颜色（使用环形采样）"""
        try:
            sample_points = self.get_sample_points()
            self.target_color = self.calculate_average_color(sample_points)

            if self.target_color == (0, 0, 0):
                self.status_label.config(text="采样失败！请调整采样参数")
                return

            self.tolerance = int(self.tolerance_entry.get() or self.tolerance)
            hex_color = '#{:02x}{:02x}{:02x}'.format(*self.target_color)
            self.target_frame.config(bg=hex_color)

            # 显示采样参数
            self.status_label.config(text=(
                f"目标颜色已锁定 | 半径: {self.ring_radius}像素 "
                f"| 采样点: {self.sample_points} "
                f"| 容差: {self.tolerance}"
            ))

            # 清空历史记录
            self.color_history.clear()

        except Exception as e:
            self.status_label.config(text=f"锁定失败: {str(e)}")

    def update_color(self):
        try:
            if not self.detecting:
                self.root.after(50, self.update_color)
                return

            # 获取采样点
            sample_points = self.get_sample_points()

            # 获取所有采样点的颜色
            point_colors = self.get_point_colors(sample_points)

            # 计算当前平均颜色
            if point_colors:
                total_r, total_g, total_b = 0, 0, 0
                for color in point_colors:
                    total_r += color[0]
                    total_g += color[1]
                    total_b += color[2]
                n = len(point_colors)
                current_color = (
                    int(total_r / n),
                    int(total_g / n),
                    int(total_b / n)
                )
            else:
                current_color = (0, 0, 0)

            # 更新当前颜色显示
            hex_color = '#{:02x}{:02x}{:02x}'.format(*current_color)
            self.current_frame.config(bg=hex_color)

            # 添加到历史记录
            if current_color != (0, 0, 0):
                self.color_history.append(current_color)

            if self.target_color and point_colors:
                # 动态调整容差
                effective_tolerance = self.update_tolerance() if self.dynamic_tolerance else self.tolerance

                # 检查每个采样点是否在容差范围内
                matched_points = 0
                for color in point_colors:
                    distance = self.calculate_color_distance(color, self.target_color)
                    if distance <= effective_tolerance:
                        matched_points += 1
                        # 只要有一个点匹配就认为检测成功
                        break

                # 检测成功
                if matched_points > 0:
                    self.root.config(bg="green")
                    self.info_label.config(text=f"★ 检测到目标颜色！匹配点: {matched_points}/{len(point_colors)} ★")
                else:
                    # 计算最小距离用于显示
                    min_distance = min(
                        self.calculate_color_distance(color, self.target_color)
                        for color in point_colors
                    )
                    self.root.config(bg="SystemButtonFace")
                    self.info_label.config(text=(
                        f"检测中... | 最小距离: {min_distance:.1f}/"
                        f"{effective_tolerance} | 按Ctrl+C退出"
                    ))

            self.root.after(50, self.update_color)

        except Exception as e:
            print(f"更新错误: {e}")
            self.root.after(50, self.update_color)

    def run(self):
        self.root.after(0, self.update_color)
        self.root.mainloop()


if __name__ == "__main__":
    detector = ColorDetector()
    try:
        detector.run()
    except KeyboardInterrupt:
        print("程序已安全退出")